import sys
import os
from threading import Thread

import numpy as np

sys.path.append('/home/cybaster/miniconda3/envs/ros1/lib/python3.8/site-packages')
import glfw
from scipy.spatial.transform import Rotation
import rospy
import rospkg
from ocs2_msgs.msg import mpc_state, mpc_input, mpc_target_trajectories, mpc_observation, mpc_flattened_controller
from ocs2_msgs.srv import reset, resetRequest
import mujoco
import mujoco.viewer


def wrap(theta) -> tuple:
    number = np.floor(theta / (2 * np.pi))
    theta -= 2 * np.pi * number
    if theta < -np.pi:
        theta += 2 * np.pi
        number -= 1
    elif theta > np.pi:
        theta -= 2 * np.pi
        number += 1
    return theta, number


class XControlAssignment:
    def __init__(self, dx: float, dy: float, ct: float, cm: float):
        super().__init__()

        self._P = np.diag([ct, ct, ct, cm])

        self._M0 = np.array([
            [1, 1, 1, 1],
            [-dx, dx, dx, -dx],
            [dy, dy, -dy, -dy],
            [1, -1, 1, -1]
        ])

        self._M0inv = np.linalg.inv(self._M0)
        self._Pinv = np.linalg.inv(self._P)

    def assign(self, taus: np.ndarray) -> np.ndarray:
        return self._M0inv @ self._Pinv @ taus


class QuadrotorSim:

    def __init__(self):
        super().__init__()
        self._robot_name = "quadrotor"
        rospy.init_node(self._robot_name + "_sim", anonymous=True)
        package_path = rospkg.RosPack().get_path("mpc_quadrotor_sim")
        xml_path = os.path.join(package_path, 'assets', 'skydio_x2', 'scene.xml')

        self._mrt_desired_frequency = 1000
        self._rate = rospy.Rate(self._mrt_desired_frequency)
        self._mpc_update_ratio = 10
        self._loop_counter = 0
        self._received = False

        self._mj_model = mujoco.MjModel.from_xml_path(xml_path)
        self._mj_data = mujoco.MjData(self._mj_model)
        self._mj_viewer = mujoco.viewer.launch_passive(self._mj_model, self._mj_data, key_callback=self._key_callback)
        self._mj_ctrl = np.zeros(4)
        self._control_assignment = XControlAssignment(0.14, 0.18, 1.0, 0.0201)

        self._state_dim = 12
        self._input_dim = 4
        self._mpc_observation_msg: mpc_observation = mpc_observation()
        self._mpc_observation_msg.time = 0.0
        self._mpc_observation_msg.state.value.extend([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
        self._mpc_observation_msg.input.value.extend([0.0, 0.0, 0.0, 0.0])
        self._mpc_observation_msg_buffer: mpc_observation = mpc_observation()
        self._mpc_flattened_controller = mpc_flattened_controller()

        self._mpc_observation_publisher = rospy.Publisher(self._robot_name + "_mpc_observation", mpc_observation,
                                                          queue_size=1)
        self._mpc_policy_subscriber = rospy.Subscriber(self._robot_name + "_mpc_policy", mpc_flattened_controller,
                                                       self._mpc_policy_callback)
        self._target_trajectories_publisher = rospy.Publisher(self._robot_name + "_mpc_target", mpc_target_trajectories,
                                                              queue_size=1)

        rospy.wait_for_service(self._robot_name + "_mpc_reset")
        self._mpc_reset_client = rospy.ServiceProxy(self._robot_name + "_mpc_reset", reset)

        self._mj_thread: Thread = Thread(target=self._mj_run, daemon=True)
        self._mj_thread.start()

    def run(self):
        rospy.spin()

    def _mpc_policy_callback(self, msg: mpc_flattened_controller):
        self._mpc_flattened_controller = msg
        self._received = True

    def _mj_run(self):

        reset_request = resetRequest()
        reset_request.reset = True
        reset_request.targetTrajectories.timeTrajectory.append(self._mpc_observation_msg.time)
        reset_request.targetTrajectories.stateTrajectory.append(self._mpc_observation_msg.state)
        reset_request.targetTrajectories.inputTrajectory.append(self._mpc_observation_msg.input)
        self._mpc_reset_client.call(reset_request)

        while not self._received:
            self._mpc_observation_publisher.publish(self._mpc_observation_msg)
            self._rate.sleep()
        self._received = False

        while self._mj_viewer.is_running():
            print("### Current time " + str(self._mpc_observation_msg.time))

            if self._loop_counter % self._mpc_update_ratio == 0:
                while not self._received:
                    pass
                self._update_input()
                self._received = False
                print("New MPC policy start at " + str(self._mpc_observation_msg.time))

            mujoco.mj_step(self._mj_model, self._mj_data)
            self._update_observation()
            self._mj_viewer.sync()

            if (self._loop_counter + 1) % self._mpc_update_ratio == 0:
                print(">>> Observation is published at " + str(self._mpc_observation_msg.time))
                self._mpc_observation_publisher.publish(self._mpc_observation_msg)

            self._loop_counter += 1
            self._rate.sleep()

    def _update_observation(self):
        pos = self._mj_data.qpos[:3]
        qu = [self._mj_data.qpos[4], self._mj_data.qpos[5], self._mj_data.qpos[6], self._mj_data.qpos[3]]
        rpy = Rotation.from_quat(qu).as_euler('xyz')
        rpy0 = list(map(wrap, self._mpc_observation_msg.state.value[3:6]))
        for i in range(3):
            if rpy[i] - rpy0[i][0] > np.pi:
                rpy[i] += (rpy0[i][1] - 1) * 2 * np.pi
            elif rpy[i] - rpy0[i][0] < -np.pi:
                rpy[i] += (rpy0[i][1] + 1) * 2 * np.pi
            else:
                rpy[i] += rpy0[i][1] * 2 * np.pi
        vel = self._mj_data.qvel[:3]
        omega = self._mj_data.qvel[3:6]
        self._mpc_observation_msg.time += 1.0 / self._mrt_desired_frequency
        self._mpc_observation_msg.state.value[:3] = pos
        self._mpc_observation_msg.state.value[3:6] = rpy
        self._mpc_observation_msg.state.value[6:9] = vel
        self._mpc_observation_msg.state.value[9:12] = omega
        self._mpc_observation_msg.input.value[:4] = self._mpc_flattened_controller.inputTrajectory[0].value

    def _update_input(self):
        self._mj_ctrl[:] = self._control_assignment.assign(self._mpc_flattened_controller.inputTrajectory[0].value)
        self._mj_data.ctrl[:] = self._mj_ctrl

    def _key_callback(self, keycode):
        mpc_target_trajectories_ = mpc_target_trajectories()
        mpc_target_trajectories_.timeTrajectory.extend(
            [self._mpc_observation_msg.time, self._mpc_observation_msg.time + 1.0])
        mpc_state_ = mpc_state()
        mpc_state_.value = [0.0] * self._state_dim
        mpc_state_.value[:6] = self._mpc_observation_msg.state.value[:6]

        if keycode == glfw.KEY_KP_1:
            mpc_state_.value[0] += 1.0
        elif keycode == glfw.KEY_KP_7:
            mpc_state_.value[0] -= 1.0
        elif keycode == glfw.KEY_KP_6:
            mpc_state_.value[1] += 1.0
        elif keycode == glfw.KEY_KP_4:
            mpc_state_.value[1] -= 1.0
        elif keycode == glfw.KEY_KP_8:
            mpc_state_.value[2] += 1.0
        elif keycode == glfw.KEY_KP_2:
            mpc_state_.value[2] -= 1.0
        elif keycode == glfw.KEY_KP_9:
            mpc_state_.value[5] += 1.0
        elif keycode == glfw.KEY_KP_3:
            mpc_state_.value[5] -= 1.0
        elif keycode == glfw.KEY_KP_5:
            pass
        else:
            return

        mpc_target_trajectories_.stateTrajectory.extend([self._mpc_observation_msg.state, mpc_state_])
        mpc_input_ = mpc_input()
        mpc_input_.value.extend([0.0] * self._input_dim)
        mpc_target_trajectories_.inputTrajectory.extend([mpc_input_, mpc_input_])
        self._target_trajectories_publisher.publish(mpc_target_trajectories_)


if __name__ == '__main__':
    node = QuadrotorSim()
    node.run()
